from ultralytics import YOLO


def predict_and_save_simple(model_path, image_dir, output_dir):
    """使用 Ultralytics 内置方法进行预测"""
    model = YOLO(model_path)
    model.eval()

    # 直接使用 predict 方法，自动处理所有细节
    results = model.predict(
        source=image_dir,
        save_txt=True,  # 保存为 txt 文件
        project=output_dir,  # 输出目录
        name=f'predictions_{model_path}',  # 子目录名称
        exist_ok=True,  # 覆盖已存在的目录
    )

    print("预测完成！结果保存在:", f"{output_dir}/predictions/labels/")


if __name__ == '__main__':
    models = ['best_10epoch.pt', 'best_100epoch.pt']
    for model in models:
        predict_and_save_simple(
            model_path=model,
            image_dir='./uav_dataset/val/images',
            output_dir='./uav_dataset/val/'
        )